import asyncio
import re

import aiosqlite
from openai import AsyncOpenAI
from tqdm.asyncio import tqdm
from itertools import cycle

# --- 1. 配置 (与您原脚本保持一致) ---
# 数据库文件路径
DB_PATH = "sentiment_analysis3.db"
# 并发请求数量
CONCURRENCY_LIMIT = 24

# LM Studio 服务器地址和模型标识符
LMSTUDIO_BASE_URL = "http://localhost:1234/v1"
MODEL_IDENTIFIERS = [
    "qwen/qwen3-4b",
    # 如果您运行了多个模型实例，可以像这样添加
    "qwen/qwen3-4b:2",
]
API_KEY = "lm-studio"

# --- 2. 新的中文系统提示 (专门用于气候相关性判断) ---
SYSTEM_PROMPT_CLIMATE = """你是一个高度精准的文本分类机器人。你的唯一任务是判断一个句子的主题是否与“气候”有关。

- 如果句子**与气候相关**，请分类为：`1`
- 如果句子**与气候无关**，请分类为：`0`
注意你的输出格式：用<result>和</result>标签对包裹你的分类结果，然后终止思考和输出
"""


async def get_sentences_to_refine(db_path: str):
    """从数据库中获取需要重新分类的句子。"""
    print(f"🔍 连接到数据库 '{db_path}' 并查找需要优化的条目...")
    async with aiosqlite.connect(db_path) as db:
        # 查询所有 result 不为 2 的记录的 id 和 sentence
        cursor = await db.execute("SELECT id, sentence FROM sentiments WHERE result != 2")
        rows = await cursor.fetchall()
        print(f"✅ 找到 {len(rows)} 条记录需要进行气候相关性二次分类。")
        return rows


async def classify_climate_relevance(client: AsyncOpenAI, sentence_id: int, sentence: str, semaphore: asyncio.Semaphore,
                                     model_identifier: str):
    """
    使用LLM判断单个句子是否与气候相关。
    返回一个元组 (sentence_id, is_climate_related)，其中 is_climate_related 为 1 (相关) 或 0 (无关)。
    """
    async with semaphore:
        try:
            response = await client.chat.completions.create(
                model=model_identifier,
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT_CLIMATE},
                    {"role": "user", "content": f'句子: "{sentence}"'}
                ],
                temperature=0.0  # 需要确定性的分类结果
            )
            content = response.choices[0].message.content.strip()
            result_match = re.search(r'<result>(.*?)</result>', content, re.DOTALL)
            if result_match:
                try:
                    result = int(result_match.group(1).strip())
                    return (sentence_id, int(result))
                except (ValueError, IndexError):
                    tqdm.write(f"⚠️ 模型在<result>标签内返回了非数字内容: '{result_match.group(1)}'")
            else:
                tqdm.write(f"⚠️ 模型未提供<result>标签。完整响应: '{content[:200]}...'")

        except Exception as e:
            tqdm.write(f"❌ 处理句子ID {sentence_id} 时发生错误 (模型: {model_identifier}): {e}")
            return (sentence_id, 1)  # 发生错误时，同样维持原样


async def main():
    """主函数，负责读取、二次分类和更新数据库。"""
    print("--- 数据库气候相关性优化脚本 ---")
    print(f"🚀 并发数: {CONCURRENCY_LIMIT}")
    print(f"� 将轮流调用模型: {MODEL_IDENTIFIERS}")
    print("-" * 40)

    # 从数据库获取需要处理的句子
    rows_to_process = await get_sentences_to_refine(DB_PATH)
    if not rows_to_process:
        print("✅ 数据库中没有需要处理的记录。脚本结束。")
        return

    client = AsyncOpenAI(base_url=LMSTUDIO_BASE_URL, api_key=API_KEY)
    semaphore = asyncio.Semaphore(CONCURRENCY_LIMIT)
    tasks = []
    model_cycler = cycle(MODEL_IDENTIFIERS)

    # 为每个需要处理的句子创建分类任务
    for row_id, sentence in rows_to_process:
        if not sentence or not sentence.strip():
            continue  # 跳过空句子

        model_id_to_use = next(model_cycler)
        task = asyncio.create_task(classify_climate_relevance(client, row_id, sentence, semaphore, model_id_to_use))
        tasks.append(task)

    updated_count = 0

    # 使用 aiosqlite 连接，在循环中处理数据库更新
    async with aiosqlite.connect(DB_PATH) as db:
        # 初始化进度条
        pbar = tqdm(total=len(rows_to_process), desc="二次分类进度", unit="句")

        # 当任务完成时立即处理
        for future in asyncio.as_completed(tasks):
            result = await future
            if result:
                sentence_id, is_climate_related = result

                # 如果模型判断句子与气候无关 (返回0)，则更新数据库
                if is_climate_related == 0:
                    try:
                        await db.execute(
                            "UPDATE sentiments SET result = 2 WHERE id = ?",
                            (sentence_id,)
                        )
                        await db.commit()
                        updated_count += 1
                        # 更新进度条的后缀信息，实时显示更新计数
                        pbar.set_postfix({"已更新为'2'的数量": updated_count})
                    except Exception as e:
                        pbar.write(f"❌ 更新数据库时发生错误 (ID: {sentence_id}): {e}")

            pbar.update(1)  # 每次有任务完成，进度条都前进一格

        pbar.close()

    print("\n" + "-" * 40)
    print("🎉 优化流程结束！")
    print(f"总共检查了 {len(rows_to_process)} 条记录。")
    print(f"其中 {updated_count} 条记录因与气候无关，其 'result' 已被更新为 2。")
    print(f"数据库 '{DB_PATH}' 已更新。")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n检测到用户中断，正在退出...")